#!/usr/bin/python
# -*- coding: utf-8 -*-

# Unit tests for comm_channels.py

import unittest
import comm_channels
import cPickle
import generic_mock

class TestCommChannel(unittest.TestCase):

    hosts_ = {}
    dispatcher_ = '127.0.0.1', 1
    hosts_[1] = '127.0.0.1', 2
    hosts_[2] = '127.0.0.1', 4
    hosts_[3] = '127.0.0.1', 5
    hosts_[4] = '127.0.0.1', 6
    pid_ = 2

    def setUp(self):
        self.comm_ = comm_channels.comm_channel(
            self.pid_, self.hosts_, self.dispatcher_)

    def testConstructor(self):
        self.assertEqual(self.comm_.dispatcher, self.dispatcher_)
        self.assertEqual(self.comm_.hosts, self.hosts_)
        self.assertEqual(self.comm_.pid, self.pid_)
        self.assertEqual(self.comm_.partners, [1, 3, 4])
        self.assertEqual(self.comm_.SEPARATOR, 'XdemX')

    def testChannelState(self):
        self.assertEqual(self.comm_.channel_state(2), False)
        self.assertEqual(self.comm_.channel_state(1), True)
        self.assertEqual(self.comm_.channel_state(3), True)
        self.assertEqual(self.comm_.channel_state(4), True)
    
    def testGetHosts(self):
        self.assertEqual(self.comm_.get_hosts(), self.hosts_)

    def testGetPID(self):
        self.assertEqual(self.comm_.get_pid(), self.pid_)

    def testNetworkSize(self):
        self.assertEqual(self.comm_.network_size(), 4)

    message = None
    destination = None
    message_sent = False
    def SendMessageStub(self, dest, msg):
        self.message = msg
        self.destination = dest
        self.message_sent = True

    def testCloseChannel(self):
        self.comm_.send_message = self.SendMessageStub
        self.comm_.close_channel(1)
        self.assertEqual(self.message, None)
        self.assertEqual(self.destination, 1)
        self.assertEqual(self.message_sent, True)


class TestTCPCommChannel(unittest.TestCase):

    hosts_ = {}
    dispatcher_ = '127.0.0.1', 1
    hosts_[1] = '127.0.0.1', 2
    hosts_[2] = '127.0.0.1', 4
    hosts_[3] = '127.0.0.1', 5
    hosts_[4] = '127.0.0.1', 6
    pid_ = 2
    
    def setUp(self):
        self.socket_ = generic_mock.GenericMock()
        self.thread_ = generic_mock.GenericMock()
        self.socketobj_ = generic_mock.GenericMock()
        self.threadobj_ = generic_mock.GenericMock()
        self.condition_ = generic_mock.GenericMock()
        self.mocks_ = [self.socket_, self.thread_, 
                       self.socketobj_, self.threadobj_,
                       self.condition_]
        self.socket_.Property('AF_INET', 'AF_INET')
        self.socket_.Property('SOCK_STREAM', 'SOCK_STREAM')
        self.socket_.StubForTesting('socket', 
                                    lambda x, y: self.socketobj_)
        self.thread_.StubForTesting('Thread', 
                                    lambda target, args: self.threadobj_)
        self.thread_.StubForTesting('Condition', 
                                    lambda: self.condition_)
        self.comm_ = comm_channels.TCP_comm_channel(
            self.pid_, self.hosts_, self.dispatcher_, 
            self.socket_, self.thread_)

    def CleanMockData(self):
        for i in self.mocks_:
            i.clear()

    def testConstructor(self):
        self.assertEqual(self.socket_, 
                         [('socket', ('AF_INET', 'SOCK_STREAM'), {}), 
                          ('socket', ('AF_INET', 'SOCK_STREAM'), {})])
        self.assertEqual([i for i, _, _ in self.thread_], 
                         ['Condition', 'Condition', 'Condition', 'Condition', 
                          'Thread', 'Thread', 'Thread',
                          'Thread', 'Thread', 'Thread'])
        self.assertEqual(self.socketobj_, 
                         [('bind', (('127.0.0.1', 4),), {}), 
                          ('listen', (10,), {}), 
                          ('connect', (('127.0.0.1', 1),), {}), 
                          ('recv', (1024,), {}), 
                          ('accept', (), {}), 
                          ('accept', (), {}), 
                          ('accept', (), {})])
        self.assertEqual(self.threadobj_, 
                         [('start', (), {}), ('start', (), {}), 
                          ('start', (), {}), ('start', (), {}), 
                          ('start', (), {}), ('start', (), {})])
        
        self.assertEqual(self.comm_.out_queue, {1: [], 3: [], 4: []})
        self.assertEqual(self.comm_.connected, {1: False, 
                                                3: False, 
                                                4: False})
        self.assertEqual(self.comm_.out_lock.keys(), [1, 3, 4])

    def testSendMessage(self):
        self.CleanMockData()
        self.comm_.send_message(3, ['dummy message'])
        self.assertEqual(self.comm_.out_queue[3], [['dummy message']])
        self.assertEqual(self.condition_, [('acquire', (), {}), 
                                           ('notify', (), {}), 
                                           ('release', (), {})])

    def testRecvMessage(self):
        self.CleanMockData()
        self.condition_.StubForTesting(
            'wait', lambda: self.comm_.in_queue.append((1, 'dummy')))
        self.assertEqual(self.comm_.recv_message(), (1, 'dummy'))
        self.assertEqual(self.condition_, [('acquire', (), {}), 
                                           ('wait', (), {}), 
                                           ('release', (), {})])
        
    def testRecvChannel(self):
        self.CleanMockData()
        sep = self.comm_.SEPARATOR
        origin = cPickle.dumps(comm_channels.origin_msg(1)) + sep
        message = cPickle.dumps('this\nis a message') + sep
        socket_recv_data = [origin, message[:5], message[5:]]
        def SocketRecvStub(unused_limit):
            if socket_recv_data:
                return socket_recv_data.pop(0) 
            else:
                return None

        channel = generic_mock.GenericMock()
        channel.StubForTesting('recv', SocketRecvStub)

        self.comm_.recv_channel(channel, 'addr')
        self.assertEqual(self.comm_.connected[1], True)
        self.assertEqual(self.comm_.in_queue, 
                         [(1, 'this\nis a message')])
        self.assertEqual(channel, [('recv', (1024,), {}), 
                                   ('recv', (1024,), {}), 
                                   ('recv', (1024,), {}), 
                                   ('recv', (1024,), {})])
  
        self.assertEqual(self.condition_, [('acquire', (), {}), 
                                           ('notify', (), {}), 
                                           ('release', (), {})])
        
    def testSendChannel(self):
        self.CleanMockData()
        sep = self.comm_.SEPARATOR
        self.comm_.out_queue[1] = ['mystring', 'oddstring']

        def SocketSendStub(object):
            if len(object) % 2 == 0:
                return len(object) / 2
            else:
                return len(object)

        self.socketobj_.StubForTesting('send', SocketSendStub)
        self.condition_.StubForTesting(
            'wait', lambda: self.comm_.out_queue[1].append(None))
                   
        self.comm_.send_channel(1)
        self.assertEqual(
            self.socketobj_, 
            [('connect', (('127.0.0.1', 2),), {}), 
             ('send', ('ccopy_reg\n_reconstructor\np1\n(ccomm_channels'
                       '\norigin_msg\np2\nc__builtin__\nobject\np3\n'
                       'NtRp4\n(dp5\nS\'pid\'\np6\nI2\nsb.XdemX',), {}),
             ('send', ('S\'mystring\'\np1\n.XdemX',), {}), 
             ('send', ('S\'oddstring\'\np1\n.XdemX',), {}), 
             ('send', ('\'\np1\n.XdemX',), {}), 
             ('close', (), {})])
        self.assertEqual(cPickle.dumps('mystring') + sep, 
                         self.socketobj_[2][1][0])
        self.assertEqual(cPickle.dumps('oddstring') + sep, 
                         self.socketobj_[3][1][0])
        self.assertEqual(
            self.condition_, 
            [('acquire', (), {}), ('release', (), {}), 
             ('acquire', (), {}), ('release', (), {}), 
             ('acquire', (), {}), ('wait', (), {}), ('release', (), {})])

if __name__ == '__main__':
    unittest.main()
